import type { H3Event } from "h3";
import { getRequestIP } from "h3";
import { RateLimiterMemory } from "rate-limiter-flexible";
import type {
  RateLimiterAbstract,
  RateLimiterRes,
} from "rate-limiter-flexible";

type RateLimitOptions = {
  /** 默认: 5000 */
  limit: number; // Limit
  /** 默认: 3500s = 1h */
  timeWindow: number; // Per second
  keyGenerator?: (e: H3Event) => string;
};

let rateLimiter: RateLimiterAbstract;

function rateLimit(config?: RateLimitOptions) {
  if (!rateLimiter) {
    rateLimiter = new RateLimiterMemory({
      points: config?.limit || 5000,
      duration: config?.timeWindow || 3600,
    });
  }
  return async (e: H3Event) => {
    const key = config?.keyGenerator
      ? config.keyGenerator(e)
      : getRequestIP(e) || "localhost";
    let res: RateLimiterRes;
    let success = true;
    try {
      res = await rateLimiter.consume(key, 1);
    } catch (e: any) {
      success = false;
      res = e;
    }
    e.res.headers.set("RateLimit-Limit", rateLimiter.points.toString());
    e.res.headers.set("RateLimit-Remaining", res.remainingPoints.toString());
    e.res.headers.set(
      "RateLimit-Reset",
      Math.ceil((Date.now() + res.msBeforeNext) / 1000).toString(),
    );
    e.res.headers.set(
      "RateLimit-Policy",
      `${rateLimiter.points};w=${rateLimiter.duration}`,
    );
    if (!success) {
      const after = Math.ceil(res.msBeforeNext / 1000);
      e.res.headers.set("Retry-After", `${after}`);
      e.res.headers.set("Content-Type", "application/json");
      e.res.status = 429;
      e.res.statusText = "Too Many Requests";
      return JSON.stringify({
        error: "rate_limit_exceeded",
        message: "You have exceeded your rate limit. Please try again later.",
        retry_after: after,
      });
    }
  };
}

export default rateLimit;
